from .binary_classifier import BinaryClassifier
from keras.models import Sequential
from keras.layers import Dense
from sklearn.preprocessing import LabelEncoder
import numpy as np
from matplotlib import pyplot as plt
from multiprocessing import Process, Manager
from sklearn.model_selection  import train_test_split


def runCategoricalClassifierModel(data):
    np.random.seed(1986)
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams['axes.unicode_minus'] = False
    
    hidden_layer_units = data["hidden_layer_units"]
    x = data["x"]
    encoder = data["encoder"]
    hidden_layer_count = data["hidden_layer_count"]
    loss = data["loss"]
    optimizer = data["optimizer"]
    train_x = data["train_x"]
    train_y = data["train_y"]
    test_x = data["test_x"]
    test_y = data["test_y"]
    epochs = data["epochs"]
    batch_size = data["batch_size"]
    chart_path = data["chart_path"]
    predict_input_values = data["predict_input_values"]
    
    # 构建模型
    model = Sequential()
    model.add(Dense(units=hidden_layer_units[0], input_dim=x.shape[1], activation='relu'))
    data["model_hidden_layers_code"] = "model.add(Dense(units=%d, input_dim=x.shape[1], activation=\'relu\'))\n" % (hidden_layer_units[0])
    for i in range(1, hidden_layer_count):
        model.add(Dense(units=hidden_layer_units[i], activation='relu'))
        data["model_hidden_layers_code"] += "model.add(Dense(units=%d, activation=\'relu\'))\n" % (hidden_layer_units[i])
    # 去掉最后一个换行字符
    data["model_hidden_layers_code"] = data["model_hidden_layers_code"][:-1] 
    model.add(Dense(units=len(encoder.classes_), activation='softmax'))
    
    # 打印模型信息
    summary_info = []
    model.summary(print_fn=lambda x:summary_info.append(x))
    data["model_summary"] = "<br/>".join(summary_info)
    
    # 编译并训练模型
    model.compile(loss=loss,
          optimizer=optimizer,
          metrics=['accuracy'])
    history = model.fit(train_x, train_y, epochs=epochs, batch_size=batch_size, verbose=0)
    
    # 绘图
    fig = plt.figure()
    ax = fig.add_subplot(111)
    # 绘制训练过程(准确度) 
    ax.plot(history.history["acc"])
    ax.set_title("模型准确度")
    ax.set_xlabel("epoch")
    ax.set_ylabel("accuracy")
    fig.savefig("%s/accuracy.png" % (chart_path))
    ax.cla()
    # 绘制训练过程(损失函数)
    ax.plot(history.history["loss"])
    ax.set_title("模型损失函数")
    ax.set_xlabel("epoch")
    ax.set_ylabel("loss_value")
    fig.savefig("%s/loss_value.png" % (chart_path))
    ax.cla()
    # 评分
    data["train_score"] = model.evaluate(train_x, train_y, verbose=0)[1]
    data["test_score"] = model.evaluate(test_x, test_y, verbose=0)[1]
    # 预测
    predict_x = np.array(predict_input_values)
    data["predict_output_values"] = encoder.inverse_transform(model.predict_classes(predict_x)).tolist()
    data["predict_output_probas"] = model.predict_proba(predict_x).tolist()


class CategoricalClassifier(BinaryClassifier):
    encoder = LabelEncoder()

    def __init__(self):
        BinaryClassifier.__init__(self)
        self.algorithm_name = "深度学习-多分类"
        self.ipynb_template_name = "categorical_classifier-template.ipynb"
        self.loss = "sparse_categorical_crossentropy"
        
    def transferXAndY(self):
        BinaryClassifier.transferXAndY(self)
        self.y = self.encoder.fit_transform(self.y)

    def implent(self):
        # 执行算法
        self.transferXAndY()
        # 拆分训练集和测试集
        (self.train_x, self.test_x, self.train_y, self.test_y) = train_test_split(self.x, self.y, train_size=self.train_size, test_size=self.test_size)
        with Manager() as manager:
            data = manager.dict()
            data["hidden_layer_units"] = self.hidden_layer_units
            data["x"] = self.x
            data["encoder"] = self.encoder
            data["hidden_layer_count"] = self.hidden_layer_count
            data["loss"] = self.loss
            data["optimizer"] = self.optimizer
            data["train_x"] = self.train_x
            data["train_y"] = self.train_y
            data["test_x"] = self.test_x
            data["test_y"] = self.test_y
            data["epochs"] = self.epochs
            data["batch_size"] = self.batch_size
            data["chart_path"] = self.chart_path
            data["predict_input_values"] = self.predict_input_values
            
            process = Process(target=runCategoricalClassifierModel, args=(data,))
            process.start()
            process.join()
            
            self.model_hidden_layers_code = data["model_hidden_layers_code"]
            self.model_summary = data["model_summary"]
            self.train_score = data["train_score"]
            self.test_score = data["test_score"]
            self.predict_output_values = data["predict_output_values"]
            self.predict_output_probas = data["predict_output_probas"]
